#!/bin/bash
# nanochat 中国环境一键配置脚本
# 适用于：Ubuntu 20.04+, CentOS 7+, macOS
# 作者：nanochat 中文社区
# 版本：1.0.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印函数
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查命令是否存在
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# 主函数
main() {
    echo ""
    echo "=========================================="
    echo "  nanochat 中国环境配置脚本"
    echo "=========================================="
    echo ""
    
    # 1. 检查系统
    print_info "检查系统环境..."
    check_system
    
    # 2. 配置 pip 镜像
    print_info "配置 pip 镜像..."
    setup_pip_mirror
    
    # 3. 配置 Rust 镜像
    print_info "配置 Rust 镜像..."
    setup_rust_mirror
    
    # 4. 配置 HuggingFace 镜像
    print_info "配置 HuggingFace 镜像..."
    setup_hf_mirror
    
    # 5. 安装 Rust（如果需要）
    if ! command_exists rustc; then
        print_info "安装 Rust..."
        install_rust
    else
        print_success "Rust 已安装: $(rustc --version)"
    fi
    
    # 6. 安装 uv（如果需要）
    if ! command_exists uv; then
        print_info "安装 uv..."
        install_uv
    else
        print_success "uv 已安装: $(uv --version)"
    fi
    
    # 7. 创建虚拟环境
    print_info "创建 Python 虚拟环境..."
    create_venv
    
    # 8. 安装依赖
    print_info "安装项目依赖..."
    install_dependencies
    
    # 9. 编译 rustbpe
    print_info "编译 rustbpe 分词器..."
    build_rustbpe
    
    # 10. 验证安装
    print_info "验证安装..."
    verify_installation
    
    # 11. 完成
    print_success "配置完成！"
    print_next_steps
}

# 检查系统
check_system() {
    # 检查操作系统
    if [[ "$OSTYPE" == "linux-gnu"* ]]; then
        OS="Linux"
    elif [[ "$OSTYPE" == "darwin"* ]]; then
        OS="macOS"
    else
        print_error "不支持的操作系统: $OSTYPE"
        exit 1
    fi
    print_success "操作系统: $OS"
    
    # 检查 Python
    if ! command_exists python3; then
        print_error "未找到 Python3，请先安装 Python 3.10+"
        exit 1
    fi
    PYTHON_VERSION=$(python3 --version | cut -d' ' -f2)
    print_success "Python 版本: $PYTHON_VERSION"
    
    # 检查 GPU
    if command_exists nvidia-smi; then
        GPU_INFO=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
        print_success "检测到 GPU: $GPU_INFO"
    else
        print_warning "未检测到 NVIDIA GPU"
    fi
}

# 配置 pip 镜像
setup_pip_mirror() {
    mkdir -p ~/.pip
    cat > ~/.pip/pip.conf << 'EOF'
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
trusted-host = pypi.tuna.tsinghua.edu.cn
timeout = 1000

[install]
trusted-host = pypi.tuna.tsinghua.edu.cn
EOF
    print_success "pip 镜像配置完成"
}

# 配置 Rust 镜像
setup_rust_mirror() {
    # 设置环境变量
    export RUSTUP_DIST_SERVER=https://rsproxy.cn
    export RUSTUP_UPDATE_ROOT=https://rsproxy.cn/rustup
    
    # 写入 shell 配置
    SHELL_RC="$HOME/.bashrc"
    if [[ "$OS" == "macOS" ]]; then
        SHELL_RC="$HOME/.zshrc"
    fi
    
    if ! grep -q "RUSTUP_DIST_SERVER" "$SHELL_RC"; then
        cat >> "$SHELL_RC" << 'EOF'

# Rust 镜像配置
export RUSTUP_DIST_SERVER=https://rsproxy.cn
export RUSTUP_UPDATE_ROOT=https://rsproxy.cn/rustup
EOF
    fi
    
    # 配置 Cargo
    mkdir -p ~/.cargo
    cat > ~/.cargo/config << 'EOF'
[source.crates-io]
replace-with = 'rsproxy-sparse'

[source.rsproxy-sparse]
registry = "sparse+https://rsproxy.cn/index/"

[net]
git-fetch-with-cli = true
EOF
    print_success "Rust 镜像配置完成"
}

# 配置 HuggingFace 镜像
setup_hf_mirror() {
    SHELL_RC="$HOME/.bashrc"
    if [[ "$OS" == "macOS" ]]; then
        SHELL_RC="$HOME/.zshrc"
    fi
    
    if ! grep -q "HF_ENDPOINT" "$SHELL_RC"; then
        echo 'export HF_ENDPOINT=https://hf-mirror.com' >> "$SHELL_RC"
    fi
    export HF_ENDPOINT=https://hf-mirror.com
    print_success "HuggingFace 镜像配置完成"
}

# 安装 Rust
install_rust() {
    curl --proto '=https' --tlsv1.2 -sSf https://rsproxy.cn/rustup-init.sh | sh -s -- -y
    source "$HOME/.cargo/env"
    print_success "Rust 安装完成: $(rustc --version)"
}

# 安装 uv
install_uv() {
    pip3 install uv -i https://pypi.tuna.tsinghua.edu.cn/simple
    print_success "uv 安装完成"
}

# 创建虚拟环境
create_venv() {
    if [ -d ".venv" ]; then
        print_warning "虚拟环境已存在，跳过创建"
    else
        uv venv
        print_success "虚拟环境创建完成"
    fi
}

# 安装依赖
install_dependencies() {
    source .venv/bin/activate
    
    # 使用 uv 安装依赖（更快）
    export UV_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
    uv sync
    
    print_success "依赖安装完成"
}

# 编译 rustbpe
build_rustbpe() {
    source .venv/bin/activate
    
    # 确保 maturin 已安装
    if ! command_exists maturin; then
        pip install maturin -i https://pypi.tuna.tsinghua.edu.cn/simple
    fi
    
    # 编译
    uv run maturin develop --release --manifest-path rustbpe/Cargo.toml
    
    print_success "rustbpe 编译完成"
}

# 验证安装
verify_installation() {
    source .venv/bin/activate
    
    echo ""
    echo "=========================================="
    echo "  环境验证"
    echo "=========================================="
    
    # Python
    echo -n "Python: "
    python --version
    
    # PyTorch
    echo -n "PyTorch: "
    python -c "import torch; print(torch.__version__)" 2>/dev/null || echo "❌ 未安装"
    
    # CUDA
    echo -n "CUDA: "
    python -c "import torch; print('✅ 可用' if torch.cuda.is_available() else '❌ 不可用')" 2>/dev/null || echo "❌ 检查失败"
    
    # Rust
    echo -n "Rust: "
    rustc --version
    
    # rustbpe
    echo -n "rustbpe: "
    python -c "import rustbpe; print('✅ 已安装')" 2>/dev/null || echo "❌ 未安装"
    
    # 磁盘空间
    echo -n "磁盘空间: "
    df -h . | tail -1 | awk '{print $4 " 可用"}'
    
    echo "=========================================="
    echo ""
}

# 打印下一步
print_next_steps() {
    echo ""
    echo "=========================================="
    echo "  🎉 安装成功！"
    echo "=========================================="
    echo ""
    echo "下一步："
    echo ""
    echo "  1. 激活虚拟环境："
    echo "     ${GREEN}source .venv/bin/activate${NC}"
    echo ""
    echo "  2. 开始训练（完整流程，约4小时）："
    echo "     ${GREEN}bash speedrun.sh${NC}"
    echo ""
    echo "  3. 或者训练小模型（快速体验）："
    echo "     ${GREEN}python scripts/base_train.py --depth=12 --device_batch_size=8${NC}"
    echo ""
    echo "  4. 查看中文教程："
    echo "     ${GREEN}cat docs/zh-CN/01-教程总览.md${NC}"
    echo ""
    echo "  5. 训练中文模型："
    echo "     ${GREEN}bash docs/zh-CN/配置工具/train_chinese.sh${NC}"
    echo ""
    echo "=========================================="
    echo ""
    echo "💡 提示："
    echo "  - 如果遇到网络问题，请检查镜像配置"
    echo "  - 训练前确保有足够的磁盘空间（至少 50GB）"
    echo "  - 单卡训练时间约为多卡的 8 倍"
    echo ""
    echo "📚 更多帮助："
    echo "  - 中文文档: docs/zh-CN/"
    echo "  - 常见问题: docs/zh-CN/03-环境配置详解.md#常见问题"
    echo "  - GitHub Issues: https://github.com/karpathy/nanochat/issues"
    echo ""
}

# 运行主函数
main "$@"

